#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <winsock2.h>
#include <ws2tcpip.h>
#include "cqupt.h"
#pragma comment(lib, "Ws2_32.lib")

#define CQUPT_EXAMPLE_SERVER_IP "127.0.0.1"

// server端处理上传请求
void cqupt_example_handle_upload(int sockfd, const cqupt_message_t* initial_message) {
    if (cqupt_validate_header(&initial_message->header) != 0) {
        fprintf(stderr, "Invalid header received.\n");
        return;
    }
    
    // 获取文件名
    char file_name[CQUPT_MAX_PAYLOAD_SIZE + 1];
    memcpy(file_name, initial_message->payload, initial_message->header.length);
    file_name[initial_message->header.length] = '\0';
    
    // 接收第二个消息：文件内容
    cqupt_message_t content_message = {0};
    if (recv(sockfd, (char *)&content_message.header, CQUPT_HEADER_SIZE, 0) <= 0) {
        fprintf(stderr, "Failed to receive file content header.\n");
        return;
    }
    
    if (cqupt_validate_header(&content_message.header) != 0) {
        fprintf(stderr, "Invalid file content header received.\n");
        return;
    }
    
    if (content_message.header.length > 0) {
        if (recv(sockfd, (char *)content_message.payload, content_message.header.length, 0) <= 0) {
            fprintf(stderr, "Failed to receive file content.\n");
            return;
        }
    }
    
    // 保存文件
    FILE *file = fopen(file_name, "wb");
    if (!file) {
        fprintf(stderr, "Failed to open file for writing: %s\n", file_name);
        return;
    }
    
    if (content_message.header.length > 0) {
        if (fwrite(content_message.payload, 1, content_message.header.length, file) != content_message.header.length) {
            fprintf(stderr, "Failed to write file data.\n");
            fclose(file);
            return;
        }
    }
    fclose(file);

    printf("File '%s' uploaded successfully.\n", file_name);
}

void cqupt_example_handle_download(int sockfd, const cqupt_message_t* initial_message) {
    if (cqupt_validate_header(&initial_message->header) != 0) {
        fprintf(stderr, "Invalid header received.\n");
        return;
    }
    
    // 获取文件名
    char file_name[CQUPT_MAX_PAYLOAD_SIZE + 1];
    memcpy(file_name, initial_message->payload, initial_message->header.length);
    file_name[initial_message->header.length] = '\0';
    
    // 尝试打开文件
    FILE *file = fopen(file_name, "rb");
    if (!file) {
        fprintf(stderr, "Failed to open file for reading: %s\n", file_name);
        // 发送失败响应
        cqupt_message_t response = {0};
        response.header.version = CQUPT_VERSION;
        response.header.type = CQUPT_ACK_FAILURE;
        response.header.length = 0;
        response.header.checksum = cqupt_calculate_checksum(&response.header, sizeof(response.header) - sizeof(response.header.checksum));
        if (send(sockfd, (char *)&response.header, CQUPT_HEADER_SIZE, 0) <= 0) {
            fprintf(stderr, "Failed to send failure response.\n");
        }
        return;
    }
    
    // 获取文件大小
    fseek(file, 0, SEEK_END);
    long file_size = ftell(file);
    fseek(file, 0, SEEK_SET);
    
    if (file_size > CQUPT_MAX_PAYLOAD_SIZE) {
        fprintf(stderr, "File too large to send: %s\n", file_name);
        fclose(file);
        // 发送失败响应
        cqupt_message_t response = {0};
        response.header.version = CQUPT_VERSION;
        response.header.type = CQUPT_ACK_FAILURE;
        response.header.length = 0;
        response.header.checksum = cqupt_calculate_checksum(&response.header, sizeof(response.header) - sizeof(response.header.checksum));
        if (send(sockfd, (char *)&response.header, CQUPT_HEADER_SIZE, 0) <= 0) {
            fprintf(stderr, "Failed to send failure response.\n");
        }
        return;
    }
    
    // 发送第一个消息：文件名确认
    cqupt_message_t message = {0};
    message.header.version = CQUPT_VERSION;
    message.header.type = CQUPT_ACK_SUCCESS;
    message.header.length = strlen(file_name);
    message.header.checksum = cqupt_calculate_checksum(&message.header, sizeof(message.header) - sizeof(message.header.checksum));
    memcpy(message.payload, file_name, strlen(file_name));
    
    if (send(sockfd, (char *)&message, CQUPT_HEADER_SIZE + strlen(file_name), 0) <= 0) {
        fprintf(stderr, "Failed to send file name confirmation.\n");
        fclose(file);
        return;
    }
    
    // 发送第二个消息：文件内容
    memset(&message, 0, sizeof(message));
    message.header.version = CQUPT_VERSION;
    message.header.type = CQUPT_CMD_DOWNLOAD;
    message.header.length = file_size;
    message.header.checksum = cqupt_calculate_checksum(&message.header, sizeof(message.header) - sizeof(message.header.checksum));
    
    if (file_size > 0) {
        if (fread(message.payload, 1, file_size, file) != file_size) {
            fprintf(stderr, "Failed to read file content.\n");
            fclose(file);
            return;
        }
    }
    fclose(file);
    
    if (send(sockfd, (char *)&message, CQUPT_HEADER_SIZE + file_size, 0) <= 0) {
        fprintf(stderr, "Failed to send file content.\n");
        return;
    }
    
    printf("File '%s' downloaded successfully.\n", file_name);
}

int main() {
    WSADATA wsaData;
    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
        fprintf(stderr, "WSAStartup failed.\n");
        return -1;
    }
    SOCKET server_socket = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
    if (server_socket == INVALID_SOCKET) {
        fprintf(stderr, "Socket creation failed.\n");
        WSACleanup();
        return -1;
    }
    struct sockaddr_in server_addr;
    server_addr.sin_family = AF_INET;
    server_addr.sin_addr.s_addr = inet_addr(CQUPT_EXAMPLE_SERVER_IP);
    server_addr.sin_port = htons(CQUPT_PORT);
    if (bind(server_socket, (struct sockaddr*)&server_addr, sizeof(server_addr)) == SOCKET_ERROR) {
        fprintf(stderr, "Bind failed.\n");
        closesocket(server_socket);
        WSACleanup();
        return -1;
    }
    if (listen(server_socket, SOMAXCONN) == SOCKET_ERROR) {
        fprintf(stderr, "Listen failed.\n");
        closesocket(server_socket);
        WSACleanup();
        return -1;
    }
    printf("Server listening on %s:%d\n", CQUPT_EXAMPLE_SERVER_IP, CQUPT_PORT);
    while (1) {
        SOCKET client_socket = accept(server_socket, NULL, NULL);
        if (client_socket == INVALID_SOCKET) {
            fprintf(stderr, "Accept failed.\n");
            continue;
        }
        printf("Client connected.\n");
        cqupt_message_t message = {0};
        
        // 接收初始消息（包含头部和负载）
        if (recv(client_socket, (char *)&message.header, CQUPT_HEADER_SIZE, 0) <= 0) {
            fprintf(stderr, "Failed to receive message header.\n");
            closesocket(client_socket);
            continue;
        }
        if (cqupt_validate_header(&message.header) != 0) {
            fprintf(stderr, "Invalid header received.\n");
            closesocket(client_socket);
            continue;
        }
        
        // 接收负载数据
        if (message.header.length > 0) {
            if (recv(client_socket, (char *)message.payload, message.header.length, 0) <= 0) {
                fprintf(stderr, "Failed to receive message payload.\n");
                closesocket(client_socket);
                continue;
            }
        }
        
        if (message.header.type == CQUPT_CMD_UPLOAD) {
            cqupt_example_handle_upload(client_socket, &message);
        } else if (message.header.type == CQUPT_CMD_DOWNLOAD) {
            cqupt_example_handle_download(client_socket, &message);
        } else {
            fprintf(stderr, "Unsupported command type: %d\n", message.header.type);
        }
        closesocket(client_socket);
        printf("Client disconnected.\n");
    }
}
